import spacy

# 加载 spaCy 模型
nlp = spacy.load("en_core_web_sm")

# 定义要读取的字符数
num_chars = 100000

print("txt start.")
# 读取文本文件的一部分
with open(r"C:\Users\86157\Desktop\course2024-2\模式识别\PatternRecognition\War and Peace.txt", "r", encoding="utf-8") as file:
    text = file.read(num_chars)

# 处理文本
doc = nlp(text)
entities = []
for ent in doc.ents:
    entities.append((ent.text, ent.label_))
print("txt done.")


# 提取关系（这里简化为提取人名和地名之间的关系）
print("start getting relations.")
relations = []
for i in range(len(doc)):
    if doc[i].ent_type_ == "PERSON":
        for j in range(i + 1, len(doc)):
            if doc[j].ent_type_ == "GPE":  # GPE: Geo-Political Entity
                relations.append((doc[i].text, doc[j].text, "from"))
print("getting relations done.")

from neo4j import GraphDatabase

# 连接到 Neo4j 数据库
uri = "bolt://localhost:7687"
user = "neo4j"
password = "密码"

print("start constructing graph.")
def create_graph(uri, user, password, entities, relations):
    driver = GraphDatabase.driver(uri, auth=(user, password))
    with driver.session() as session:
        # 创建实体
        for entity in entities:
            session.run("MERGE (n:Entity {name: $name, type: $type})", name=entity[0], type=entity[1])

        # 创建关系
        for relation in relations:
            session.run("MATCH (a:Entity {name: $name1}), (b:Entity {name: $name2}) "
                        "MERGE (a)-[r:RELATED_TO {type: $type}]->(b)", name1=relation[0], name2=relation[1],
                        type=relation[2])

    driver.close()


create_graph(uri, user, password, entities, relations)

print("all done.")